import random
import re

def clean_text(text):
    """
    对传入的文本进行清洗，可以根据需要加入更多规则
    """
    # # 移除HTML标签（如果存在）
    # text = re.sub(r'<[^>]+>', '', text)
    # 移除多余的空白字符
    text = re.sub(r'\s+', ' ', text)
    return text.strip()

def stream_lines(file_path):
    """
    按行生成文件内容，同时进行清洗
    """
    with open(file_path, 'r', encoding='utf-8') as f:
        for line in f:
            cleaned = clean_text(line)
            # 如果清洗后的结果为空，则跳过
            if cleaned:
                yield cleaned

def mixed_stream(file1, file2, buffer_size=10000, output_file='combined_cleaned_data.txt'):
    # 分别创建两个生成器
    gen1 = stream_lines(file1)
    gen2 = stream_lines(file2)

    # 开启输出文件
    with open(output_file, 'w', encoding='utf-8') as out_f:
        buffer = []
        exhausted1 = exhausted2 = False

        while True:
            # 随机选择从哪个生成器取数据
            source_choice = random.choice([1, 2])
            if source_choice == 1 and not exhausted1:
                try:
                    buffer.append(next(gen1))
                except StopIteration:
                    exhausted1 = True
            elif source_choice == 2 and not exhausted2:
                try:
                    buffer.append(next(gen2))
                except StopIteration:
                    exhausted2 = True

            # 当两个数据源都用完或缓冲区达到阈值时，进行批次内打乱和写入
            if (len(buffer) >= buffer_size) or (exhausted1 and exhausted2 and buffer):
                random.shuffle(buffer)
                for line in buffer:
                    out_f.write(line + "\n")
                buffer = []

            # 如果两个生成器都已耗尽，则退出循环
            if exhausted1 or exhausted2:
                break

# 调用函数：调整文件路径和缓冲大小参数以适合实际情况
mixed_stream('data/raw/books.txt', 'data/raw/wikipedia.txt', buffer_size=10000, output_file='data/combined.txt')
